﻿// FindPartition.cpp : 定义静态库的函数。
//

#include "pch.h"
#include "framework.h"

#include "Q:\\includes\\ESP.h"
#include <corecrt_io.h>
#pragma comment(lib , "Q:\\build\\lib\\lib\\static\\ESP.lib")

#include "Q:\includes\ThrowError_CPP.h"
#include <algorithm>

std::vector<char> GetVolumeLetterB();

#include <iostream>
#include <string>
#include <sstream>

std::string FP_Char2Str(char c);

std::tuple<MainProgramPartitionLetter, ReinstallPartitionLetter, EfiPartitionLetter> RiProgram_Common_GetPartition(
	_Out_ ErrorCode Status
) {
	
	MainProgramPartitionLetter mainProgramPartitionLetter = NULL;
	ReinstallPartitionLetter reinstallPartitionLetter = NULL;
	EfiPartitionLetter efiPartitionLetter = NULL;

	//挂载EFI分区
	//先卸载，再挂载
	std::string EFI = "";
	try {
		Esp_UnmountEfiPartitionString();
		std::string Letter = Esp_MountEfiPartitionString();
		if (Letter == "") {
			throw;
		}

		//没有EFI分区，报错返回
		if (Letter == "") {
			Status = __RIPROG_NO_EFI_PARTITION;
			return { NULL,NULL,NULL };
		}
	}
	catch(const char* err){
		//失败，返回
		Status = __RIPROG_GET_EFI_PARTITION_FAIL;
		return { NULL,NULL,NULL };
	}

	std::vector<char> AllLetter;
	//挂载成功，开始获取所有盘符
	try {
		AllLetter = GetVolumeLetterB();
	}
	catch (const char* err) {
		//失败，返回
		Status = __RIPROG_GET_PARTITION_FAIL;
		return { NULL,NULL,NULL };
	}

	//没有分区，返回
	if (AllLetter.size() == 0) {
		Status = __RIPROG_NO_PARTITION;
		return { NULL,NULL,NULL };
	}

	//遍历判断文件是否存在
	for (int i = 0; i < AllLetter.size(); i++) {
		char Letter = AllLetter[i];

		//拼接路径
		std::string MainProgramPath = FP_Char2Str(Letter)+ ":\\AZPEP_SysRi__MainProgramPartition";
		std::string ReInstallPartition = FP_Char2Str(Letter) + ":\\AZPEP_SysRi_Tool__ReInstallPartition";
		std::string EfiPartition = FP_Char2Str(Letter) + ":\\AZPEP_SysRi_Tool__EfiPartition";

		//获取文件
		//如果文件不存在，或有重复文件，则报错返回

		//AZPEP_SysRi__MainProgramPartition
		if (_access(MainProgramPath.c_str(), 0) !=-1  && mainProgramPartitionLetter == NULL) mainProgramPartitionLetter = Letter;
		else if (mainProgramPartitionLetter != NULL) { Status = __RIPROG_TOO_MANY_MAINPROG_LETTER; return { NULL,NULL,NULL }; }

		//AZPEP_SysRi__MainProgramPartition
		if (_access(ReInstallPartition.c_str(), 0) != -1 && reinstallPartitionLetter == NULL) reinstallPartitionLetter = Letter;
		else if (reinstallPartitionLetter != NULL) { Status = __RIPROG_TOO_MANY_SYSTEM_PARTITION; return { NULL,NULL,NULL }; }

		//AZPEP_SysRi__MainProgramPartition
		if (_access(EfiPartition.c_str(), 0) != -1 && efiPartitionLetter == NULL)  efiPartitionLetter = Letter;
		else if (efiPartitionLetter != NULL) { Status = __RIPROG_TOO_MANY_EFI_PARTITION; return { NULL,NULL,NULL }; }

	}

	//如果某项没设置，则报错返回
	if (mainProgramPartitionLetter == NULL) { Status = __RIPROG_NO_MAINPROG_PARTITION_LETTER; return { NULL,NULL,NULL }; }
	if (reinstallPartitionLetter == NULL) { Status = __RIPROG_NO_SYSTEM_PARTITION_LETTER; return { NULL,NULL,NULL }; }
	if (efiPartitionLetter == NULL) { Status = __RIPROG_NO_EFI_PARTITION_LETTER; return { NULL,NULL,NULL }; }

	Status = SUCCESS;
	return { mainProgramPartitionLetter,reinstallPartitionLetter,efiPartitionLetter };
}

char GetEfiPartitionLetter(_Out_ ErrorCode Status) {
	char efiPartitionLetter =NULL;
	
	//挂载EFI分区
	//先卸载，再挂载
	try {

		Esp_UnmountEfiPartitionString();
		std::string EFI = Esp_MountEfiPartitionString();
		if (EFI == "") {
			throw;
		}

		//没有EFI分区，报错返回
		if (EFI == "") {
			Status = __RIPROG_NO_EFI_PARTITION;
			return NULL;
		}
		
		for (int i = 0; i < EFI.size(); i++) {
			std::string EfiPartition = FP_Char2Str(EFI[i]) + ":\\AZPEP_SysRi_Tool__EfiPartition";
			//AZPEP_SysRi__MainProgramPartition
			if (_access(EfiPartition.c_str(), 0) != -1 && efiPartitionLetter == NULL) { efiPartitionLetter = EFI[i]; break; }
			else if (efiPartitionLetter != NULL) { Status = __RIPROG_TOO_MANY_EFI_PARTITION; return NULL; }
		}

		if (efiPartitionLetter == NULL) { Status = __RIPROG_NO_EFI_PARTITION_LETTER; return NULL; }

		Status = SUCCESS;
		return efiPartitionLetter;
	}

	catch (const char* err) {
		//失败，返回
		Status = __RIPROG_GET_EFI_PARTITION_FAIL;
		return NULL;
	}

}
char GetMainProgramLetter(_Out_ ErrorCode Status) {
	char mainProgramPartitionLetter = NULL;
	std::vector<char> AllLetter;

	try {
		AllLetter = GetVolumeLetterB();
	}
	catch (const char* err) {
		//失败，返回
		Status = __RIPROG_GET_PARTITION_FAIL;
		return NULL;
	}

	for (int i = 0; i < AllLetter.size(); i++) {
		char Letter = AllLetter[i];

		std::string MainProgramPath = FP_Char2Str(Letter) + ":\\AZPEP_SysRi__MainProgramPartition";
		if (_access(MainProgramPath.c_str(), 0) != -1 && mainProgramPartitionLetter == NULL) { mainProgramPartitionLetter = Letter; break; }
		else if (mainProgramPartitionLetter != NULL) { Status = __RIPROG_TOO_MANY_MAINPROG_LETTER; return NULL; }
	}

	if (mainProgramPartitionLetter == NULL) { Status = __RIPROG_NO_EFI_PARTITION_LETTER; return NULL; }

	Status = SUCCESS;
	return mainProgramPartitionLetter;
}
char GetReinstallPartitionLetter(_Out_ ErrorCode Status) {
	char ReInstallPartitionLetter = NULL;
	std::vector<char> AllLetter;

	try {
		AllLetter = GetVolumeLetterB();
	}
	catch (const char* err) {
		//失败，返回
		Status = __RIPROG_GET_PARTITION_FAIL;
		return NULL;
	}

	for (int i = 0; i < AllLetter.size(); i++) {
		char Letter = AllLetter[i];

		std::string MainProgramPath = FP_Char2Str(Letter) + ":\\AZPEP_SysRi_Tool__ReInstallPartition";
		if (_access(MainProgramPath.c_str(), 0) != -1 && ReInstallPartitionLetter == NULL) { ReInstallPartitionLetter = Letter; break; }

		else if (ReInstallPartitionLetter != NULL) { Status = __RIPROG_TOO_MANY_MAINPROG_LETTER; return NULL; break; }
	}

	if (ReInstallPartitionLetter == NULL) { Status = __RIPROG_NO_EFI_PARTITION_LETTER; return NULL; }

	//std::cout << "\n" << ReInstallPartitionLetter << "\n" << std::endl;

	Status = SUCCESS;
	return ReInstallPartitionLetter;
}

/// <summary>
///  获取磁盘上所有已使用的盘符
/// </summary>
/// <returns></returns>
std::vector<char> GetVolumeLetterB() {
	std::vector<char> DriveLetter;

	//创建完全盘符列表
	std::vector<char> AllLetter;
	for (char i = 'A'; i <= 'Z'; i++) {
		std::string LetterPath;
		LetterPath += i;
		LetterPath += ":\\";

		//当不等于=1是，说明分区已挂载。将其添加到 DriveLetter
		if (_access(LetterPath.c_str(), 0) != -1) {
			DriveLetter.push_back(i);
		}
	}

	//如果没有多于盘符，报错ERROR_GETESPINFO_MO_MORE_AVAIBLE_LETTER，返回
	if (DriveLetter.size() == 0) {
		throw;
	}

	return DriveLetter;
}

std::string FP_Char2Str(char c) {
	std::string s(1, c);
	return s;
}